# -*- coding: utf-8 -*-
"""
Created on Sat Nov 18 19:36:23 2017

@author: Luther
"""

import numpy as np
import matplotlib.pyplot as plt
import tushare as ts

r1 = 0.1
r2 = 0.2
sig1 = 0.1
sig2 = 0.3
corr = -0.5
return_list = []
risk_list = []
for w in np.arange(0.001, 1, 0.001):
    ret = w * r1 + (1 - w) * r2
    risk = (w**2) * (sig1**2) + (
        (1 - w)**2) * (sig2**2) + 2 * w * (1 - w) * sig1 * sig2 * corr
    return_list.append(ret)
    risk_list.append(np.sqrt(risk))

plt.plot(risk_list, return_list)
plt.show()

szzs = ts.get_k_data(
    '000001', start='1996-01-01', end='2016-12-31', index=True)
s = szzs.close
lns = np.log(s)
lns = lns[::250]
s = lns.diff()
sd = s.std()
mean = s.mean()
print('mean:{0:.2f}%, sd:{1:.2f}%'.format(100 * mean, 100 * sd))

rf = 0.04
rm = mean
sigf = 0
sigm = sd
y_l = []
x_l = []
for w in np.arange(0.001, 1, 0.001):
    y = (1 - w) * rf + w * rm
    x = (w**2) * (sigm**2)
    y_l.append(y)
    x_l.append(np.sqrt(x))

plt.plot(x_l, y_l)
plt.axis([0, 0.35, 0.03, 0.1])
plt.show

p = 10
mu = 0.02
sig = 0.1
p_list = []
p_list.append(p)
rand = np.random.randn(90)
for i in range(len(rand)):
    growth = mu + sig * rand[i]
    p = p * (1 + growth)
    p_list.append(p)

plt.plot(p_list, 'r')
